import '@assets/css/app.css';
import { $, vec3, sphere, camera, plane, phong, union, checker, ray3 } from "@utils/functions";
import { Sphere, Camera, Color, Union, Plane, Ray3 } from "@script/index";

function render_depth(canvas, scene: Union, camera: Camera, max_depth: number) {
  const context = canvas.getContext('2d');
  const { width, height }: any = canvas.attributes;
  const w = width.value;
  const h = height.value;
  context.fillStyle = 'rgb(0, 0, 0)';
  context.fillRect(0, 0, w, h);
  const imageData = context.getImageData(0, 0, w, h);
  const pixels = imageData.data;

  scene.initialize();
  camera.initialize();

  let i = 0;
  for (let y = 0; y < h; ++y) {
    let sy = 1 - y / h;
    for (let x = 0; x < w; ++x) {
      let sx = x / w;
      let ray = camera.generateRay(sx, sy);
      let result = scene.intersect(ray);
      if (result.geometry) {
        let depth = 255 - Math.min(result.distance / max_depth * 255, 255);
        pixels[i] = depth;
        pixels[i + 1] = depth;
        pixels[i + 2] = depth;
        pixels[i + 3] = 255;
      }
      i += 4;
    }
  }
  context.putImageData(imageData, 0, 0);
}

function render_normal(canvas, scene: Union, camera: Camera) {
  const context = canvas.getContext('2d');
  const { width, height } = canvas.attributes;
  const w = width.value;
  const h = height.value;
  context.fillStyle = 'rgb(0, 0, 0)';
  context.fillRect(0, 0, w, h);
  const imageData = context.getImageData(0, 0, w, h);
  const pixels = imageData.data;

  scene.initialize();
  camera.initialize();

  let i = 0;
  for (let y = 0; y < h; ++y) {
    let sy = 1 - y / h;
    for (let x = 0; x < w; ++x) {
      let sx = x / w;
      let ray = camera.generateRay(sx, sy);
      let result = scene.intersect(ray);
      if (result.geometry) {
        pixels[i] = (result.normal.x + 1) * 128;
        pixels[i + 1] = (result.normal.y + 1) * 128;
        pixels[i + 2] = (result.normal.z + 1) * 128;
        pixels[i + 3] = 255;
      }
      i += 4;
    }
  }
  context.putImageData(imageData, 0, 0);
}

function ray_trace(canvas, scene: Union, camera: Camera) {
  const context = canvas.getContext('2d');
  const { width, height } = canvas.attributes;
  const w = width.value;
  const h = height.value;
  context.fillStyle = 'rgb(0, 0, 0)';
  context.fillRect(0, 0, w, h);
  const imageData = context.getImageData(0, 0, w, h);
  const pixels = imageData.data;

  scene.initialize();
  camera.initialize();

  let i = 0;
  for (let y = 0; y < h; ++y) {
    let sy = 1 - y / h;
    for (let x = 0; x < w; ++x) {
      let sx = x / w;
      let ray = camera.generateRay(sx, sy);
      let result = scene.intersect(ray);
      if (result.geometry) {
        let color = result.geometry.material.sample(ray, result.position, result.normal);
        pixels[i] = color.r * 255;
        pixels[i + 1] = color.g * 255;
        pixels[i + 2] = color.b * 255;
        pixels[i + 3] = 255;
      }
      i += 4;
    }
  }
  context.putImageData(imageData, 0, 0);
}

function ray_trace_reflection(canvas, scene, camera, max_reflect: number) {
  const context = canvas.getContext('2d');
  const { width, height }: any = canvas.attributes;
  const w = width.value;
  const h = height.value;
  context.fillStyle = 'rgb(0, 0, 0)';
  context.fillRect(0, 0, w, h);
  const imageData = context.getImageData(0, 0, w, h);
  const pixels = imageData.data;

  scene.initialize();
  camera.initialize();

  let i = 0;
  for (let y = 0; y < h; ++y) {
    let sy = 1 - y / h;
    for (let x = 0; x < w; ++x) {
      let sx = x / w;
      let ray = camera.generateRay(sx, sy);
      let color = reflection_recursive(scene, ray, max_reflect);
      pixels[i++] = color.r * 255;
      pixels[i++] = color.g * 255;
      pixels[i++] = color.b * 255;
      pixels[i++] = 255;
    }
  }
  context.putImageData(imageData, 0, 0);
}

function reflection_recursive(scene: Union, ray: Ray3, max_reflect: number): Color {
  let result = scene.intersect(ray);
  if (result.geometry) {
    let reflectiveness = result.geometry.material.reflectiveness;
    let color = result.geometry.material.sample(ray, result.position, result.normal);
    color = color.multiply(1 - reflectiveness);
    if (reflectiveness > 0 && max_reflect > 0) {
      let r = result.normal.multiply(-2 * result.normal.dot(ray.direction)).add(ray.direction);
      ray = ray3(result.position, r);
      let reflected_color = reflection_recursive(scene, ray, max_reflect - 1);
      color = color.add(reflected_color.multiply(reflectiveness));
    }
    return color;
  } else {

    return Color.black;
  }
}

function _main() {
  const canvas0 = $('#first-canvas');
  render_depth(
    canvas0,
    union([
      plane(vec3(0, 1, 0), 0),
      sphere(vec3(0, 10, -10), 10)
    ]),
    camera(
      vec3(0, 10, 10),
      vec3(0, 0, -1),
      vec3(0, 1, 0),
      90
    ),
    20);

  const canvas1 = $('#second-canvas');
  render_normal(
    canvas1,
    union([
      plane(vec3(0, 1, 0), 0),
      sphere(vec3(0, 10, -10), 10)
    ]),
    camera(
      vec3(0, 10, 10),
      vec3(0, 0, -1),
      vec3(0, 1, 0),
      90
    )
  );

  const canvas2 = $('#third-canvas');

  const plane1: Plane = plane(vec3(0, 1, 0), 0);
  const sphere1: Sphere = sphere(vec3(-10, 10, -10), 10);
  const sphere2: Sphere = sphere(vec3(10, 10, -10), 10);
  plane1.material = checker(0.1, 0.5);
  sphere1.material = phong(Color.red, Color.white, 16, 0.25);
  sphere2.material = phong(Color.blue, Color.white, 16, 0.25);

  ray_trace(
    canvas2,
    union([plane1, sphere1, sphere2]),
    camera(vec3(0, 5, 15), vec3(0, 0, -1), vec3(0, 1, 0), 90)
  );

  const canvas3 = $('#forth-canvas');
  ray_trace_reflection(
    canvas3,
    union([plane1, sphere1, sphere2]),
    camera(vec3(0, 5, 15), vec3(0, 0, -1), vec3(0, 1, 0), 90),
    3
  );

}

(function (): void {
  _main();
})();
